# -*- coding: utf-8 -*-
# @Author  : XiYuan
# @Time    : 2023/12/4 21:40
import concurrent.futures
import queue
from time import sleep

# 创建配置对象（默认从 ini 文件中读取配置）
from DrissionPage._configs.chromium_options import ChromiumOptions
from DrissionPage._pages.web_page import WebPage


class TabPool:

    def __init__(self, size=3, mode='d', timeout=None, chromium_options=None, session_or_options=None):
        self._thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=size)
        self._tab_pool = queue.Queue()
        self._tab_dic = {}
        # 创建配置对象（默认从 ini 文件中读取配置）
        co = ChromiumOptions()
        # 设置不加载图片、静音
        co.mute().set_load_mode('eager')
        self._webpage = WebPage(mode, timeout, chromium_options, session_or_options)
        self._create_web_tab_pool(size)

    def _create_web_tab_pool(self, size=5):
        for i in range(0, size + 1):
            self._create_tab()

    def _create_tab(self):
        tab = self._webpage.new_tab()
        tab = self._webpage.get_tab(tab)
        self._tab_pool.put(tab)
        self._tab_dic[tab] = 0

    def get_active_web_tab(self):
        tab = self._tab_pool.get()  # 随机获取其中一个tab
        count = self._tab_dic.get(tab)
        self._tab_dic[tab] = count + 1
        # self._time_to_return_tab(tab)
        return tab

    def close_web_page_pool(self):
        """
        确保任务都已经完成才能调用这个方法
        """
        self._thread_pool.shutdown(wait=True)
        while not self._tab_pool.empty():
            self._tab_pool.get().close()

    def return_active_web_tab(self, tab):
        """
        使用完进行归还到tab池中
        """
        if tab.states.is_loading:
            sleep(1)  # 休息一秒在归还
            self.return_active_web_tab(tab)
        else:
            tab.run_cdp("Page.stopLoading")
            # 使用次数上限为100次，超过就关闭掉,这是为了解决tab请求过多后，会阻塞请求的bug
            if self._tab_dic.get(tab) > 100:
                tab.close()  # 关闭标签页
                self._create_tab()
            else:
                self._tab_pool.put(tab)

    def operate(self, func, **kwargs):
        """
        :param func: 要执行的方法
                kwargs:额外的参数
        :return:
        """

        def carry_out_func():
            tab = self.get_active_web_tab()
            func(tab, kwargs)
            self.return_active_web_tab(tab)

        self._thread_pool.submit(carry_out_func)


def aaa(tab, kwargs):
    print(kwargs)
    tab.get("https://www.baidu.com")
    tab.ele('#kw').input('hello world')
    tab.ele('#su').click()


if __name__ == '__main__':
    pool = TabPool(size=8)
    for i in range(1, 10000):
        pool.operate(aaa, kwargs=123)
    pool.close_web_page_pool()

